import pickle
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, activations, losses, optimizers, metrics, callbacks
import numpy as np
import os
import sys
import matplotlib.pyplot as plt

FEATURE_PICKLE_PATH = r'_save\trans_learn_tf1x_on_tf2x_1of3_extract_feature.py\v7.0\bottleneck.pickle'
MODEL_PATH = r'_save\trans_learn_tf2x_2of3_train_model.py\v3.1\trans_learn_self_model.dat'
IMG_ROOT_DIR = '../../../../../large_data/CV2/_many_files/flower_photos_liuqilong/'


def sep(label = '', cnt=32):
    print('-' * cnt, label, '-' * cnt, sep='')


tf.random.set_seed(1)
np.random.seed(1)

sep('Load features')
with open(FEATURE_PICKLE_PATH, 'br') as f:
    pickle_data = pickle.load(f)

idx2label = pickle_data['idx2label']
label2idx = pickle_data['label2idx']
print(idx2label)
print(label2idx)
N_CLS = len(idx2label.keys())

x_train = pickle_data['x_train'].reshape(-1, 2048)
y_train = pickle_data['y_train'].reshape(-1, 1)
print('x_train', x_train.shape)
print('y_train', y_train.shape)

x_test = pickle_data['x_test'].reshape(-1, 2048)
y_test = pickle_data['y_test'].reshape(-1, 1)
print('x_test', x_test.shape)
print('y_test', y_test.shape)

x_val = pickle_data['x_val'].reshape(-1, 2048)
y_val = pickle_data['y_val'].reshape(-1, 1)
print('x_val', x_val.shape)
print('y_val', y_val.shape)

x_path_test = pickle_data['x_path_test']

sep('Load moel')
print('Loading model ...')
model = keras.models.load_model(MODEL_PATH)
print('Loaded model')
model.summary()

sep('predict')
spr = 4
spc = 4
spn = 0
cnt = spr * spc
print('Predicting ...')
pred = model.predict(x_test[:cnt])
pred = pred.argmax(axis=1)
print('pred', pred.shape)
print('Predicted')

sep('plotting')
plt.figure(figsize=[14, 6])
for i in range(cnt):
    relative_path = x_path_test[i]
    print(i, relative_path)
    path = os.path.join(IMG_ROOT_DIR, relative_path)
    y = y_test[i][0]
    p = pred[i]
    res = 'V' if y == p else 'X'
    title = f'{idx2label[y]}:{idx2label[p]}({res})'

    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title(title)
    plt.axis('off')
    img = plt.imread(path)
    plt.imshow(img)

plt.show()
